import rrd_base
import re
import time
import snmp

import color

try:
    from pyrrd.rrd import DataSource, RRA, RRD
    from pyrrd.graph import DEF, CDEF, VDEF
    from pyrrd.graph import LINE, AREA, GPRINT
    from pyrrd.graph import ColorAttributes, Graph
except:
    print "Couldn't load required modules. You probably need to install it with:"
    sys.exit('sudo apt-get install python-pyrrd python-pysnmp4')

class RrdLoadavg(rrd_base.RrdBase):
    def __init__(self, domain):
        super(RrdLoadavg,self).__init__(domain=domain, group="cpu", instance="loadavg")

    def create_file(self):
        dataSources = []
        dataSources.append(DataSource( dsName='loadavg', dsType='GAUGE', heartbeat=self.heartbeat()))
        dataSources.append(DataSource( dsName='n_cpus',  dsType='GAUGE', heartbeat=self.heartbeat()))

        roundRobinArchives = []
        for func in ['AVERAGE', 'MAX', 'MIN']:
            for step,row in self.rra_pattern():
                roundRobinArchives.append(RRA(cf=func, xff=0.5, steps=step, rows=row))

        self.myRRD = RRD(self.filename, ds=dataSources, rra=roundRobinArchives, step=self.data_step())
        self.myRRD.create()

    def graph(self):
        d_avg    = DEF(rrdfile=self.myRRD.filename, vname='a', cdef="AVERAGE", dsName='loadavg')
        n_cpus   = DEF(rrdfile=self.myRRD.filename, vname='b', cdef="AVERAGE", dsName='n_cpus')
        d_min    = DEF(rrdfile=self.myRRD.filename, vname='c', cdef="MIN",     dsName='loadavg')
        d_max    = DEF(rrdfile=self.myRRD.filename, vname='d', cdef="MAX",     dsName='loadavg')
        dn_avg   = CDEF(vname="loadavg_normalized",   rpn="{0},{1},/".format(d_avg.vname, n_cpus.vname))
        dn_max   = CDEF(vname="max_normalized",       rpn="{0},{1},/".format(d_max.vname, n_cpus.vname))
        dn_min   = CDEF(vname="min_normalized",       rpn="{0},{1},/".format(d_min.vname, n_cpus.vname))
        vn_peak  = VDEF(vname="v0",                   rpn="{0},MAXIMUM".format(dn_max.vname))
        vn_nadir = VDEF(vname="v1",                   rpn="{0},MINIMUM".format(dn_min.vname))
        vn_avg   = VDEF(vname="v2",                   rpn="{0},AVERAGE".format(dn_avg.vname))
        area_0 = AREA(   defObj=dn_avg,   color=color.RGBA['green'], legend='Average')
        line_0 = LINE(1, defObj=dn_max,   color=color.RGBA['red'],   legend="Maximum")
        line_1 = LINE(1, defObj=dn_min,   color=color.RGBA['blue'],  legend="Minumum")
        line_2 = LINE(1, defObj=vn_peak,  color=color.RGBA['red'],   legend="Peak")
        line_3 = LINE(1, defObj=vn_nadir, color=color.RGBA['blue'],  legend="Nadir")
        line_4 = LINE(1, defObj=vn_avg,   color=color.RGBA['green'])

        for tf in self.graph_time_frames().keys():
            g = Graph(self.image_name(tf),
                      start=int(time.time())-self.graph_time_frames()[tf], end=int(time.time()),
                      lower_limit=0.00001,
                      title='"Load Average for {0} - Updated {1}"'.format(self.data_domain, time.strftime("%Y%m%d at %H:%M")),
                      height=self.graph_height, width=self.graph_width,
                      x_grid=self.grid[tf],
                      step=self.graph_step_times[tf])
            g.data.extend([d_avg, n_cpus, dn_avg, area_0, d_max, dn_max, d_min, dn_min, vn_peak, vn_nadir, vn_avg, line_2, line_3, line_4])
            if self.graph_time_frames()[tf] > self.graph_time_frames()["1day"]:
                g.data.extend([line_0])
            g.write()

            if False:
                g = Graph(self.icon_name(tf),
                          start=int(time.time())-self.graph_time_frames()[tf], end=int(time.time()),
                          lower_limit=0.00001,
                          no_legend=True,
                          height=self.icon_height, width=self.icon_width,
                          step=self.graph_step_times[tf])
                g.data.extend([d_avg, n_cpus, dn_avg, area_0])
                g.write()

class RrdLoadavgLocal(RrdLoadavg):
    def __init__(self, domain):
        super(RrdLoadavgLocal,self).__init__(domain=domain)

    def update(self):
        f = open("/proc/loadavg")
        for line in f:
            words = line.split()

        f = open("/proc/cpuinfo")
        n_cpu = 0
        for line in f:
            if re.match("[pP]rocessor", line):
                n_cpu += 1

        self.myRRD.bufferValue(int(round(time.time())), words[1], n_cpu)
        self.myRRD.update()

class RrdLoadavgSnmp(RrdLoadavg):
    def __init__(self, domain):
        super(RrdLoadavgSnmp,self).__init__(domain=domain)

    def update(self):
        load = snmp.get((1,3,6,1,4,1,2021,10,1,5,2))
        self.myRRD.bufferValue(int(round(time.time())), float(load), 1)
        self.myRRD.update()
